import numpy as np
import scipy
from scipy.signal import butter, filtfilt
from scipy.io import wavfile
from mne.io import read_raw_edf
import os
import librosa
from pynwb import NWBHDF5IO
import pandas as pd

def z_score_normalization(data):
    mean = np.mean(data, axis=0, keepdims=True)
    std = np.std(data, axis=0, keepdims=True)
    return (data - mean) / std

# 从实信号转换为解析信号，解析信号的幅值表示瞬时振幅。通过计算瞬时振幅，了解不同时间点的强度变化
hilbert3 = lambda x: scipy.signal.hilbert(x, scipy.fftpack.next_fast_len(len(x)),axis=0)[:len(x)]
def extractHG(data, sr, windowLength, class_start, dir_path, file):
    # linear detrend
    data = scipy.signal.detrend(data, axis=0) # 去除线性趋势
    numWindows = int(np.floor(data.shape[0]/(winL*sr)))

    # Create feature space  窗口平均
    data = np.abs(hilbert3(data))
    class_name = class_start
    for win in range(numWindows):
        start = int(np.floor(win * sr))
        stop = int(np.floor(start + windowLength * sr))
        feat = data[start:stop, :]
        feat = z_score_normalization(feat)
        np.save(os.path.join(dir_path, file[:-4]+'_'+str(class_name)), feat)
        class_name += 1


if __name__ == '__main__':
    # 划窗，窗口长度为1s
    winL = 1
    root_data_path = '/root/data/video_decoding'
    eeg_path = os.path.join(root_data_path, 'cut_data_zscore')      # original file path
    outputs = os.path.join(root_data_path, 'npy_data_class_30_zscore')   # output feature

    eeg_sr = 400
    time_files = os.listdir(eeg_path)
    for data_id in ['train', 'val']:
        file_names = os.listdir(os.path.join(eeg_path, data_id))
        dir_path = f'{outputs}/{data_id}'
        for file in file_names:
            print(file)
            if file[16]=='1':
                class_start = 0
            elif file[16]=='2':
                class_start = 15
            # neural data
            eeg_data = np.load(os.path.join(eeg_path, data_id, file))  # read data
            eeg_data = eeg_data.transpose((1, 0))
            # extract HG feature     window average
            feat = extractHG(eeg_data, eeg_sr, winL, class_start, dir_path, file)
            # np.save(os.path.join(f'./npy_data/{data_id}/{file}'), feat)
